#!/usr/bin/env python
# coding: utf-8

# # Wilson Loop

# These codes are used to load wilson loop data and show wilson loop.

# In[4]:


import numpy as np
import matplotlib.pyplot as plt


# In[5]:


def jackknife(data,jaxies=0):
    nconf = data.shape[jaxies]
    jdata = (np.sum(data,jaxies)[np.newaxis,...]-data)/(nconf-1)
    aver = np.average(data,jaxies)
    error = np.std(jdata,jaxies) * (nconf-1)**0.5
    return jdata,aver,error
def effectivemass(data,confaxies=0,axies=1):
    nconf = data.shape[confaxies]
    tmp = -np.log(np.roll(data,shift=-1,axis=axies) / data)
    aver = np.average(tmp,confaxies)
    error = np.std(tmp,confaxies) * (nconf -1)**0.5
    return tmp,aver,error
def showwilsonloop(aver,error,plotr=None,filename=None,**params):
    plt.figure()
    if "title" in params.keys():
        plt.title(params["title"])
    else:
        plt.title("Wilson loop")
    if "yscale" in params.keys():
        plt.yscale(params["yscale"])
    if "ylim" in params.keys():
        plt.ylim(params["ylim"])
    if "xlim" in params.keys():
        plt.xlim(params["xlim"])
    if "xlabel" in params.keys():
        plt.xlabel(params["xlabel"])
    if "ylabel" in params.keys():
        plt.ylabel(params["ylabel"])
    trange = np.arange(aver.shape[0])
    for i in plotr:
        plt.errorbar(trange,aver[trange,i],error[trange,i],label="r = %i"%i)
    plt.legend()
    if filename is not None:
        plt.savefig(filename)
    else:
        plt.show()
    


# In[6]:


data = np.load("./purgauge_2432_WilsonLoop.npz")
wilsonloop1,wilsonloop2,wilsonloop3 = data["arr_0"],data["arr_1"],data["arr_2"]


# * wilson loop 1 for space-like planar

# In[7]:


wilsonloop1_j,aver_wilsonloop1,error_wilsonloop1 = jackknife(wilsonloop1)
showwilsonloop(aver_wilsonloop1,error_wilsonloop1,plotr=np.arange(5),yscale="log",xlabel="$\hat{t}$",ylabel="w(t,r)",title="space-like wilson loop",filename="wilsonloop1.pdf")


# In[8]:


wilsonloop1_mass,aver_wilsonloop1_mass,error_wilsonloop1_mass = effectivemass(wilsonloop1_j)
showwilsonloop(aver_wilsonloop1_mass,error_wilsonloop1_mass,plotr=np.arange(5),xlim=(0,10),ylim=(0,2),xlabel="$\hat{t}$",ylabel="$\hat{V}(t,r)$",title="space-like wilson loop",filename="wilsonloop1_V.pdf")


# * wilson loop 2 for time-like planar 

# In[9]:


wilsonloop2_j,aver_wilsonloop2,error_wilsonloop2 = jackknife(wilsonloop2)
showwilsonloop(aver_wilsonloop2,error_wilsonloop2,plotr=np.arange(5),yscale="log",xlabel="$\hat{t}$",ylabel="w(t,r)",title="time-like wilson loop",filename="wilsonloop2.pdf")


# In[10]:


wilsonloop2_mass,aver_wilsonloop2_mass,error_wilsonloop2_mass = effectivemass(wilsonloop2_j)
showwilsonloop(aver_wilsonloop2_mass,error_wilsonloop2_mass,plotr=np.arange(5),xlim=(0,10),ylim=(0,2),xlabel="$\hat{t}$",ylabel="$\hat{V}(t,r)$",title="time-like wilson loop",filename="wilsonloop2_V.pdf")


# * wilson loop 3 for off-axis time-like. the first nr is for $\sqrt{2}$ off-axis and the next nr/2 is for $\sqrt{5}$ off-axis and the rest nr is for $\sqrt{3}$ off-axis.

# In[11]:


wilsonloop3_j,aver_wilsonloop3,error_wilsonloop3 = jackknife(wilsonloop3)
showwilsonloop(aver_wilsonloop3,error_wilsonloop3,plotr=np.arange(5),yscale="log",xlabel="$\hat{t}$",ylabel="w(t,r)",title="off-axis time-like wilson loop",filename="wilsonloop3.pdf")


# In[12]:


wilsonloop3_mass,aver_wilsonloop3_mass,error_wilsonloop3_mass = effectivemass(wilsonloop3_j)
showwilsonloop(aver_wilsonloop3_mass,error_wilsonloop3_mass,plotr=np.arange(5),xlim=(0,10),ylim=(0,2.5),xlabel="$\hat{t}$",ylabel="$\hat{V}(t,r)$",title="off-axis time-like wilson loop 3",filename="wilsonloop3_V.pdf")


# * show the potential $V(r)$

# In[13]:


nr = 24
choicet=3
Vr = np.array([aver_wilsonloop2_mass[choicet,0],aver_wilsonloop3_mass[choicet,0],aver_wilsonloop3_mass[choicet,nr+nr//2],aver_wilsonloop2_mass[choicet,1],aver_wilsonloop3_mass[choicet,nr],aver_wilsonloop3_mass[choicet,1],aver_wilsonloop2_mass[choicet,2],aver_wilsonloop3_mass[choicet,nr+nr//2 + 1],aver_wilsonloop2_mass[choicet,3]])
Vr_err = np.array([error_wilsonloop2_mass[choicet,0],error_wilsonloop3_mass[choicet,0],error_wilsonloop3_mass[choicet,nr+nr//2],error_wilsonloop2_mass[choicet,1],error_wilsonloop3_mass[choicet,nr],error_wilsonloop3_mass[choicet,1],error_wilsonloop2_mass[choicet,2],error_wilsonloop3_mass[choicet,nr+nr//2 + 1],error_wilsonloop2_mass[choicet,3]])
r = np.sort(np.array([np.arange(1,nr),np.arange(1,nr)*2**0.5,np.arange(1,nr)*3**0.5,np.arange(1,nr)*5**0.5]).flatten())


# In[14]:


plt.figure()
lent=len(Vr)
plt.title("static potential")
plt.xlabel("$\hat{r}$")
plt.ylabel("$\hat{V}(r)$")
plt.ylim((0,1))
plt.xlim((0,5))
plt.minorticks_on()
plt.tick_params(direction='in', length=3, width=0.5, colors='k',
               grid_color='r', grid_alpha=0.5,axis="both",which="both",top=True,right=True)
plt.tick_params(length=5, width=1,which="major")
#plt.grid()
plt.errorbar(r[:lent],Vr,Vr_err,label="V(r)",fmt="o")
plt.legend()
plt.savefig("Vr.pdf")


# * fit the static potential

# In[15]:


#import lsqfit
#import gvar


# In[17]:


#wilsonloop2.shape,wilsonloop3.shape


# In[ ]:




